package datamining.algorithms.gspan;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Iterator;

import datamining.algorithms.gspan.AttackGraphEdge.Type;

import utils.Database;

public class AttackGraphGenerator {

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		// TODO Auto-generated method stub
		Database db = new Database("com.mysql.jdbc.Driver", "jdbc:mysql://127.0.0.1:3306/snort", "root","rootadmin");
		String[] queryList = new String[5];
		
		// ALL
		queryList[0] = "SELECT TIMESTAMP,plugin_sid,p.name,INET_NTOA(ip_src) AS src_ip,INET_NTOA(ip_dst)AS dst_ip "
				+ "FROM test.`acid_event_copy` e JOIN ossim.plugin_sid p ON (e.plugin_id=p.plugin_id AND e.plugin_sid=p.sid) "
				+ " ORDER BY TIMESTAMP;";
		//DARPA99
		queryList[1] = "SELECT TIMESTAMP,plugin_sid,p.name,INET_NTOA(ip_src) AS src_ip,INET_NTOA(ip_dst)AS dst_ip "
				+ "FROM test.`acid_event_copy` e JOIN ossim.plugin_sid p ON (e.plugin_id=p.plugin_id AND e.plugin_sid=p.sid) "
				+ "WHERE e.plugin_id=1001 and UNIX_TIMESTAMP(timestamp)<UNIX_TIMESTAMP(\"2012-10-29 21:14:53\")"
				+ " ORDER BY TIMESTAMP;";
		//DARPA2000 LLS_DDOS_1.0-inside 
		queryList[2] = "SELECT TIMESTAMP,plugin_sid,p.name,INET_NTOA(ip_src) AS src_ip,INET_NTOA(ip_dst)AS dst_ip "
				+ "FROM test.`acid_event_copy` e JOIN ossim.plugin_sid p ON (e.plugin_id=p.plugin_id AND e.plugin_sid=p.sid)"
				+ " WHERE e.plugin_id=1001 and UNIX_TIMESTAMP(timestamp)>UNIX_TIMESTAMP(\"2012-10-29 21:14:53\") and "
				+ "UNIX_TIMESTAMP(timestamp)<=UNIX_TIMESTAMP(\"2012-10-30 00:20:58\")"
				+ "ORDER BY TIMESTAMP;";
		//DARPA2000 LLS_DDOS_2.0-inside 
		queryList[3] = "SELECT TIMESTAMP,plugin_sid,p.name,INET_NTOA(ip_src) AS src_ip,INET_NTOA(ip_dst)AS dst_ip "
				+ "FROM test.`acid_event_copy` e JOIN ossim.plugin_sid p ON (e.plugin_id=p.plugin_id AND e.plugin_sid=p.sid)"
				+ " WHERE e.plugin_id=1001 and UNIX_TIMESTAMP(timestamp)>UNIX_TIMESTAMP(\"2012-10-30 00:20:58\") "
				+ "ORDER BY TIMESTAMP;";
		
		BufferedReader br=null;
		String queryString = "";
		String fileNamePrefix = "conf/data/";
		int choice = 0;
		int maxnode = 1000;
		int maxEdge = 1000;

		try {
			String input="";
			br = new BufferedReader(new InputStreamReader(System.in));
			do {
				System.out.println("请输入数据集编号:（0默认全部） \n" + "1. DARPPA99 第一个月数据\n"
						+ "2. DARPA2000 LLS_DDOS_1.0-inside 场景数据\n"
						+ "3. DARPA2000 LLS_DDOS_2.0-inside 场景数据");

				input = br.readLine();
				if (!input.equals(""))
					choice = Integer.parseInt(input);
			}while (choice < 0 || choice > 3);
			
			queryString = queryList[choice];
			
			switch (choice) {
			case 0:
				fileNamePrefix += "ALL";
				break;
			case 1:
				fileNamePrefix += "DARPA99";
				break;
			case 2:
				fileNamePrefix += "DARPA2000-L1";
				break;
			case 3:
				fileNamePrefix += "DARPA2000-L2";
				break;
			default:
				fileNamePrefix += "";
			}
			
			do {
				System.out.println("请输入最大节点数限制: 默认1000，0为不限制");
				input = br.readLine();
				if (!input.equals(""))
					maxnode = Integer.parseInt(input);
			} while (maxnode < 0);

			do {
				System.out.println("请输入最大边数限制: 默认1000，0为不限制");
				input = br.readLine();
				if (!input.equals(""))
					maxEdge = Integer.parseInt(input);
			} while (maxEdge < 0);

		} catch (NumberFormatException e1) {
			// TODO Auto-generated catch block
			e1.printStackTrace();
		} catch (IOException e1) {
			// TODO Auto-generated catch block
			e1.printStackTrace();
		} 
		
		try{
			
			ResultSet rs = db
					.database_execute_single_command(queryString);

			ArrayList<AttackGraphEdge> currentEdges = new ArrayList<AttackGraphEdge>();	//当前缓冲区内的边
			ArrayList<AttackGraphEdge> ALLEdges = new ArrayList<AttackGraphEdge>();	//所有的边
			
			boolean matched[] = new boolean[6];	//0:src-dst-label; 1:src-dst; 2:src; 3:dst; 4:a.src = b.dst; 5:a.dst = b.src
			
			int edgeCount = 0;
			int idCount = 0;
			for (edgeCount = 0, idCount = 0; rs.next() && (maxnode == 0 ? true
					: (idCount < maxnode)) && (maxEdge == 0 ? true
							: (ALLEdges.size() < maxEdge)); ++edgeCount) {
				
				String src_ip = rs.getString("src_ip");
				String dst_ip = rs.getString("dst_ip");
				String plugin_sid = rs.getString("plugin_sid");
				
				boolean done= false;
				Iterator<AttackGraphEdge> it = currentEdges.iterator();
				
				int idA = 0;
				int idB = 1;
				String src_tmp = "";
				String dst_tmp = "";
				String labeltmp = "";
				
				AttackGraphEdge[] matchedEdges = new AttackGraphEdge[6];
				for(int i=0;i<6;i++)	//初始化
					matched[i] = false;
				
				//将准备加入的新边与已存在所有的边进行匹配，以决定新边的加入位置
				while (it.hasNext()) {
					AttackGraphEdge edgeTmp = it.next();
					src_tmp = edgeTmp.nodeA.label;
					dst_tmp = edgeTmp.nodeB.label;
					labeltmp = edgeTmp.label;
					idA = edgeTmp.nodeA.id;
					idB = edgeTmp.nodeB.id;
					
					if(src_ip.equals(src_tmp) && dst_ip.equals(dst_tmp) && edgeTmp.nodeB.outDegree == 0){ //源目地址相同: A->B, A->B
						if (labeltmp.equals(plugin_sid) && !matched[0]) { 	//事件类型也相同
							matched[0] = true;
							break;
						}else if(! labeltmp.equals(plugin_sid) && ! matched[1]){	//事件类型不同
							matched[1] = true;
							matchedEdges[1] = edgeTmp;
							it.remove();
							continue;
						}
					}else if (edgeTmp.nodeA.isEnable() && dst_ip.equals(src_tmp) && ! matched[5]) { // 新事件目标ip与已有事件源ip一致: A->B, (C->A)
						matched[5] = true;
						matchedEdges[5] = edgeTmp;
						//FIXME:去掉失效边
						edgeTmp.nodeA.disable();
						if(!edgeTmp.nodeB.isEnable())
							it.remove();
						break;
					}else if (edgeTmp.nodeA.isEnable() && src_ip.equals(src_tmp) && !dst_ip.equals(dst_tmp) && ! matched[2]) { // 新事件源ip与已有事件源ip一致，目标不一致: A->B, A->C
						matched[2] = true;
						matchedEdges[2] = edgeTmp;
						continue;
					}else if(edgeTmp.nodeB.isEnable() && !src_ip.equals(src_tmp) && dst_ip.equals(dst_tmp) && ! matched[3]){ //新事件目标ip与已有事件目标ip一致，源不一致:: A->B, C->B
						if(edgeTmp.nodeB.outDegree == 0){
							matched[3] = true;
							matchedEdges[3] = edgeTmp;
						}else
							edgeTmp.nodeB.disable();
						
						continue;
					}else if(edgeTmp.nodeB.isEnable() && src_ip.equals(dst_tmp)){	//新事件源与已有事件目标相同:: A->B, (B->C)
						//it.remove();
						
						if(! matched[4]){
							matched[4] = true;
							matchedEdges[4] = edgeTmp;
						}
					}
				}
				
				
				
				AttackGraphEdge edge = null;
				
				//确定新边的加入位置
				if(!matched[0]){
					if(matched[1]){	//源目地址相同: A->B, A->B, 但事件类型不同
						idA = ++idCount;
						idB = matchedEdges[1].nodeB.id;
						AttackGraphNode newNode =  new AttackGraphNode(src_ip, idA);
						System.out.println("被匹配的边的类型：" + matchedEdges[1].type);
						if(matchedEdges[1].type != Type.edge_end && matchedEdges[1].type != Type.edge_start && matchedEdges[1].type != Type.edge_mid)
							matchedEdges[1].setType(Type.edge_start);
						else if(matchedEdges[1].type == Type.edge_end || matchedEdges[1].type == Type.edge_start)
							matchedEdges[1].setType(Type.edge_mid);
						
						//FIXME: 注意测试
						//insertString = newEdge("", idA, dst_ip, idB, plugin_sid);
						edge = new AttackGraphEdge(newNode, matchedEdges[1].nodeB, plugin_sid, Type.edge_end);
						matchedEdges[1].nodeB = newNode;
					}else if(matched[2] && ! matched[3]){	// 存在事件与新事件源一致、目标不一致，不存在事件与新事件目标一致: A->B, (A->C)
						idA = matchedEdges[2].nodeA.id;
						idB = ++idCount;
						//insertString = srcExtEdge(idA, dst_ip, idB, plugin_sid);
						edge = new AttackGraphEdge(matchedEdges[2].nodeA, dst_ip,idB, plugin_sid, Type.src);
					}else if(! matched[2] && matched[3]){	//存在事件与新事件目标一致、源不一致，不存在事件与新事件源一致、目标不一致: A->B, (C->B)
						idA = ++idCount;
						idB = matchedEdges[3].nodeB.id;
						edge = new AttackGraphEdge(src_ip, idA, matchedEdges[3].nodeB, plugin_sid, Type.dst);
					}else if(matched[2] && matched[3]){	//存在事件与新事件目标一致、源不一致，也存在事件与新事件源一致、目标不一致: A->B, C->D, (A->D)
						idA = matchedEdges[2].nodeA.id;
						idB = matchedEdges[3].nodeB.id;
						edge = new AttackGraphEdge(matchedEdges[2].nodeA, matchedEdges[3].nodeB, plugin_sid, Type.src_dst);
					}else if(matched[4]){	//新事件源与已有事件目标相同:: A->B, (B->C)
						idA = matchedEdges[4].nodeB.id;
						idB = ++idCount;
						//insertString = srcExtEdge(idA, dst_ip, idB, plugin_sid);
						edge = new AttackGraphEdge(matchedEdges[4].nodeB, dst_ip,idB, plugin_sid, Type.src);
					}else{	//全新的事件或新事件目标与已有事件源一致
						idA = ++idCount;
						idB = ++idCount;
						//insertString = newEdge(src_ip, idA, dst_ip, idB, plugin_sid);
						edge = new AttackGraphEdge(src_ip, idA, dst_ip,idB, plugin_sid, Type.none);
					}
					
					currentEdges.add(edge);	//存储新键入的边到临时匹配集合
					ALLEdges.add(edge);	//存储新加入的边到全部边的集合
				}
				
				System.out.println("plugin_sid = " + plugin_sid
						+ ", src_ip = " + src_ip
						+ ", dst_ip = " + dst_ip);
			}
			
			String outFileName = fileNamePrefix
					+ ("-" + idCount + "nodes")
					+ ("-" + ALLEdges.size() + "edges")
					+ ".dot";
			
			writeDotFile(ALLEdges, outFileName);	//写Dot文件
			
		} catch (SQLException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
	}
	
	/**
	 * 将以当前id作为起点的边全部移除
	 * @param idA
	 * @param list
	 */
	public static void removeEdgeOut(int idA, ArrayList<AttackGraphEdge> list){
		Iterator<AttackGraphEdge> it = list.iterator();
		
		while(it.hasNext()){
			AttackGraphEdge edgeTmp = it.next();
			
			if(edgeTmp.nodeA.id == idA)
				it.remove();
		}
	}
	/**
	 * 将所有边转换为Dot文件
	 * @param list
	 */
	public static void writeDotFile(ArrayList<AttackGraphEdge> list, String outFileName) {
		System.out.println("-----------------------------------------------------------");
		System.out.println("writing result to  file: " + outFileName);
		System.out.println("-----------------------------------------------------------");
		try {
			FileWriter fw;
			fw = new FileWriter(outFileName);
			BufferedWriter bw = new BufferedWriter(fw);
			bw.write("digraph G {\n");

			Iterator<AttackGraphEdge> it = list.iterator();

			while (it.hasNext()) {
				AttackGraphEdge edgeTmp = it.next();
				bw.write(edgeTmp.toString());
			}
			bw.write("}");
			bw.flush();
			bw.close();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}

	/**
	 * 移除链表中失效的边
	 * @param list
	 * @return
	 */
	public static int removeDisabledEdge(ArrayList<AttackGraphEdge> list){
		int count = 0;
		Iterator<AttackGraphEdge> it = list.iterator();

		while (it.hasNext()) {
			AttackGraphEdge edgeTmp = it.next();
			if(!edgeTmp.nodeA.isEnable() && !edgeTmp.nodeB.isEnable()){
				it.remove();
				count ++;
			}
		}
		return count;
	}
	
}
